import numpy as np
import cv2


def extractComponent(self, image, label_image, label):
    component = np.zeros(image.shape, np.uint8)
    component[label_image == label] = image[label_image == label]
    return component

img = cv2.imread('IMG2.jpeg')
img = cv2.resize(img,(600,800))
Z = img.reshape((-1, 3))

# convert to np.float32

Z = np.float32(Z)

  # define criteria, number of clusters(K) and apply kmeans()

criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)

K = 6

ret, label, center = cv2.kmeans(Z, K, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)

  # Now convert back into uint8, and make original image

center = np.uint8(center)

res = center[label.flatten()]

res2 = res.reshape((img.shape))

lable_index = label.reshape((img.shape[0],img.shape[1]))

# extracted = extractComponent(img,res2,2)

component = np.zeros(img.shape, np.uint8)
component[lable_index == 4] = img[lable_index == 4]



cv2.imshow('res2', component)

cv2.waitKey(0)

cv2.destroyAllWindows()